# nostrademus.py
#
# Nostrademus e um sistema de previsao de cotacoes para o mercado de acoes
# em Python (2.6.5) desenvolvido para a disciplina Introducao a Inteligencia
# Artificial
#
# Software Dependencies:
#     PyEvolve 0.6rc1
#     Matplotlib 0.99.1.1
#     PyDot 1.0.25
#
# November 2011, UNICAMP, Brazil
#
# Authors:
# Celso Almeida (locatelli@gmail.com)
# Gabriell Orisaka (orisaka@gmail.com)
# Jorge Pizzolatto Filho (jorge.pzt@gmail.com)
# Jose Lucas (jlucasns@gmail.com)

import glob
import math
import optparse
import os
import random
import sys

try:
    from pyevolve import GTree
    from pyevolve import GSimpleGA
    from pyevolve import Selectors
    from pyevolve import Consts
except:
    print 'PyEvolve is required!'
    sys.exit(-1)

try:
    import matplotlib.pyplot as plt
except:
    PYPLOT = False
else:
    PYPLOT = True

try:
    import pydot
except:
    PYDOT = False
else:
    PYDOT = True


# Global stuff
data = []
options = None
plabel = 'fech'


# Functions available to non-terminal nodes
@GTree.gpdec(representation='+', color='lawngreen')
def gp_add(a, b): return a + b

@GTree.gpdec(representation='-', color='lawngreen')
def gp_sub(a, b): return a - b

@GTree.gpdec(representation='avg', color='lawngreen')
def gp_avg(a, b): return (a + b) / 2.0

@GTree.gpdec(representation='max', color='lawngreen')
def gp_max(a, b): return max(a, b)

@GTree.gpdec(representation='min', color='lawngreen')
def gp_min(a, b): return min(a, b)

@GTree.gpdec(representation='iflt', color='lawngreen')
def gp_iflt(a, b, c, d):
    if a < b:
        return c
    else:
        return d

@GTree.gpdec(representation='ifgt', color='lawngreen')
def gp_ifgt(a, b, c, d):
    if a > b:
        return c
    else:
        return d

@GTree.gpdec(representation='and', color='lawngreen')
def gp_and(a, b, c, d):
    if a and b:
        return c
    else:
        return d

@GTree.gpdec(representation='or', color='lawngreen')
def gp_or(a, b, c, d):
    if a or b:
        return c
    else:
        return d


# The fitness function
def fitness_func(genome):
    global data

    code_comp = genome.getCompiledCode()

    error = 0.0
    for i in range(len(data) - options.n - 1, len(data) - 1):
        forecast = eval(code_comp)
        error += abs(forecast - data[i + 1][plabel])

    return error


# Print the fittest individual of every 10 iterations (pydot is implicitly
# required by pyevolve)
def step_callback(gp_engine):
    global options

    gen = gp_engine.getCurrentGeneration()
    if gen % 10 == 0:
        filename = '{stock}_{gen}.png'.format(stock=options.stock, gen=gen)
        try:
            GTree.GTreeGP.writePopulationDot(gp_engine, options.output + '/' + filename, 'png', 0, 1)
        except Exception as err:
            #print err
            pass


# Main genetic programming procedure
def run_gp():
    global options

    genome = GTree.GTreeGP()
    genome.setParams(max_depth=options.mdepth, method='ramped')
    genome.evaluator.set(fitness_func)

    ga = GSimpleGA.GSimpleGA(genome)
    
    # Set terminal and non-terminal nodes
    terminals = ''
    for i in xrange(options.n):
        args = []
        for label in data[0].keys():
            args += [i, label]
        terminals += 'data[i-%d][\'%s\'] ' * len(data[0].keys()) % tuple(args)
    terminals = terminals.split()
    ga.setParams(gp_terminals = terminals + ['ephemeral:random.randint(0, 100)'],
                 gp_function_prefix = 'gp')

    # Set some gp parameters
    print 'GP Parameters:'
    print '\tGenerations:', options.gen
    print '\tCrossover Rate:', options.cvr
    print '\tMutation Rate:', options.mur
    print '\tPopulation size:', options.pop
    print ''
    ga.setMinimax(Consts.minimaxType['minimize'])
    ga.selector.set(Selectors.GRouletteWheel)
    ga.setGenerations(options.gen)
    ga.setCrossoverRate(options.cvr)
    ga.setMutationRate(options.mur)
    ga.setPopulationSize(options.pop)
    if PYDOT:
        ga.stepCallback.set(step_callback)

    # Start iterations
    print 'Iterating...'
    if options.verbose:
        freq_stats = 10
    else:
        freq_stats = 0
    ga(freq_stats=freq_stats)

    # Get best individual
    best = ga.bestIndividual()
    best_code = best.getCompiledCode()

    # Get prediction
    i = len(data) - 1
    prediction = eval(best_code)
    prel = (prediction / data[i][plabel]) - 1

    # Hard decision
    hard_decision = 'SELL'
    if prediction >= data[-1][plabel]:
        hard_decision = 'BUY'

    # Calculate accuracy
    hits = 0
    for i in range(len(data) - options.et - 1, len(data) - 1):
        pred = eval(best_code)
        if (data[i + 1][plabel] - data[i][plabel]) >= 0 and (pred - data[i][plabel]) >= 0:
            hits += 1
        elif (data[i + 1][plabel] - data[i][plabel]) < 0 and (pred - data[i][plabel]) < 0:
            hits += 1
    accuracy = float(hits) / options.et

    print ''
    print 'Best individual:'
    print '\tScore:', best.score
    print '\tExpression:', best.getPreOrderExpression()
    print '\tPredition:', prediction, '{0:.2}% (Current: {cur})'.format(prel * 100, cur=data[-1][plabel])
    print '\tHard decision:', hard_decision
    print '\tAccuracy:', accuracy
    print ''

    if options.results:
        fp = open(options.results, 'a')
        fp.write('%s,%.2f,%.2f,%s,%.2f\n' % (options.stock, prediction, prel * 100, hard_decision, accuracy * 100))
        fp.close()

    if PYDOT:
        try:
            GTree.GTreeGP.writePopulationDot(ga, options.output + '/' + options.stock + '_best.png', 'png', 0, 1)
        except Exception as err:
            #print err
            pass

    if PYPLOT:
        rprice = []
        rdate = []
        odate = range(len(data) - options.n - 1, len(data))
        for i in odate:
            rprice.append(data[i][plabel])
            rdate.append(i - len(data) + 1)

        pprice = []
        pdate = []
        for i in range(len(data) - options.et - 1, len(data)):
            pprice.append(eval(best_code))
            pdate.append(i - len(data) + 2)

        plt.figure()
        plt.plot(rdate, rprice, 'b-', pdate, pprice, 'ro')
        plt.xlim(-options.n, 2)
        plt.ylim(min(rprice) * 0.8, max(rprice) * 1.2)
        plt.xlabel('day')
        plt.ylabel('price')
        plt.legend(('previous prices','predicted'))
        plt.title('Nostrademus: ' + options.stock)
        plt.savefig(options.output + '/' + options.stock + '_chart.png')
        if options.verbose:
            plt.show()


# Load stock data
def load_data(filename):
    global data

    file = open(filename)

    labels = file.readline().strip().split(',')
    
    data = []
    for line in file:
        values = line.strip().split(',')
        values = map(float, values)
        values = dict(zip(labels, values))

        if values.has_key('date'):
            del values['date']

        data.append(values)

    file.close()


# Command line entry-point
def main():
    global data
    global options

    parser = optparse.OptionParser()
    parser.add_option('-d', '--data-dir', dest='data_dir', help='Directory with stock data. Default: data', default='data')
    parser.add_option('-o', '--output-dir', dest='output', help='Output directory. Default: .', default='.')
    parser.add_option('-s', '--stock', dest='stock', help='Stock name')
    parser.add_option('-n', '--observation-time', dest='n', type='int', default=10, help='Observation time. Default: 10')
    parser.add_option('-e', '--evaluation-time', dest='et', type='int', default=10, help='Evaluation time. Default: 10')
    parser.add_option('-p', '--population', dest='pop', type='int', default=50, help='Population size [1 - inf]. Default: 50')
    parser.add_option('-g', '--generations', dest='gen', type='int', default=100, help='Number of generations [1 - inf]. Default: 100')
    parser.add_option('-c', '--cross-over-rate', dest='cvr', type='float', default=1.00, help='Crossover rate [0.00 - 1.00]. Default: 1.00')
    parser.add_option('-m', '--mutation-rate', dest='mur', type='float', default=0.25, help='Mutation rate [0.00 - 1.00]. Default: 0.25')
    parser.add_option('-x', '--max-tree-depth', dest='mdepth', type='int', default=5, help='Max tree depth [1 - inf]. Default: 5')
    parser.add_option('-r', '--results', dest='results', help='Result file. Default: None', default='')
    parser.add_option('-l', '--list', action='store_true', dest='list', help='List available stocks and exit')
    parser.add_option('-v', '--verbose', action='store_true', dest='verbose', help='Verbose mode')
    options, args = parser.parse_args()

    if options.list:
        for file in sorted(glob.glob(options.data_dir + '/*.csv')):
            print os.path.basename(file).split('.')[0]
        sys.exit(0)

    if not (options.stock and options.n):
        print 'Not enough arguments. Use -h or --help for help.'
        sys.exit(-1)

    if options.pop <= 1 or options.gen <= 1 or options.mdepth <= 1:
        print 'Invalid range for parameter! Use -h or --help for help.'
        sys.exit(-1)

    if options.cvr < 0.0 or options.cvr > 1.0:
        print 'Invalid range for parameter! Use -h or --help for help.'
        sys.exit(-1)

    if options.mur < 0 or options.mur > 1.0:
        print 'Invalid range for parameter! Use -h or --help for help.'
        sys.exit(-1)

    if options.stock != '*':
        stock_file = options.data_dir + '/' + options.stock + '.csv'

        if not glob.glob(stock_file):
            print 'File ' + stock_file + ' not found!'
            sys.exit(-1)

        if options.output != '.':
            if not os.access(options.output, os.F_OK):
                os.mkdir(options.output)

        load_data(stock_file)

        if options.n > 0.5 * len(data):
            print 'Observation time greater than stock data available!'
            sys.exit(-1)

        run_gp()
    else:
        for stock_file in sorted(glob.glob(options.data_dir + '/*.csv')):
            options.stock = os.path.basename(stock_file).split('.')[0]
            print options.stock

            if options.output != '.':
                if not os.access(options.output, os.F_OK):
                    os.mkdir(options.output)

            load_data(stock_file)

            if options.n > 0.5 * len(data):
                print 'Observation time greater than stock data available for', options.stock
                continue

            run_gp()
    
    return 0


if __name__ == '__main__':
    main()

